from yacf.Frontend.C99 import c99_ast as c_ast

import sys


from yacf.Backends.Common.Writers.CommonWriter import OffsetNodeVisitor
from yacf.Frontend.SymbolTable import getSymbolTableFromNode, IdentifierNotFound
from yacf.Backends.Common.TemplateEngine.Functions import is_complex_type
from yacf.MiddleEnd.Loop.Analysis import ParametrizeLoopTool, NonOmpCompilant

class FlopCountVisitor(OffsetNodeVisitor):
    """ Generates FlopCount code from the IR

    """
    inside = False

    def __init__(self, contableVars = [], st = None):
        self._contableVars = contableVars
        self._flop = 0
        self._memAccess = 0
        self._currentAccess = 0 # Temp. var to store partial count results
        self._flopCount_subtree = c_ast.Constant(type = "integer", value = "0")
        self._memAccess_subtree = c_ast.Constant(type = "integer", value = "0")
        self._ifCount = 0

        if st:
            self._st = st
        self._st = None

    def _lookUpAndGet(self,node):
        if not self._st:
            self._st = getSymbolTableFromNode(node)
        try:
            return self._st.lookUp(node)
        except IdentifierNotFound:
            return None

    def getFlopCount(self,):
        return max(self._flop,1)

    def getMemAccess(self,):
        return max(self._memAccess,1)


    def getFlopCountSubtree(self,):
        return self._flopCount_subtree

    def getMemAccessSubtree(self,):
        return self._memAccess_subtree

    def getIfCount(self,):
        return self._ifCount


    # ********** Visit **********

    def visit_For(self, node, offset = 0):
        """ A for loop increments the number of memAccess and flops by
            the number of iterations. """
        nIter_subtree = c_ast.Constant(type = "integer", value = "1");
        try:
            la = ParametrizeLoopTool()
            lp = la.parametrize(node)
            nIter_subtree = lp['number_of_iterations']
        except NonOmpCompilant:
            print "Non canonical loop ignored"
        self.visit(node.stmt)
        self._memAccess_subtree = c_ast.BinaryOp(op = '*', 
                                                  left = nIter_subtree,
                                                  right = self._memAccess_subtree)
        self._flopCount_subtree = c_ast.BinaryOp(op = '*', 
                                                  left = nIter_subtree,
                                                  right = self._flopCount_subtree)
       

    def visit_BinaryOp(self, node, offset = 0):
        self._currentAccess = 0
        self.constant = 0
        self.visit(node.left)
        self.visit(node.right)
        #if self._currentAccess > 0:
        if not node.op in ['<','>','<=','>=','&&','||'] and not self.constant:
            # print "Op is : " + str(node)
            self._flop += 1
            self._flopCount_subtree = c_ast.BinaryOp( op = "+", 
                                                     left = c_ast.Constant(type = "integer", value = "1"),
                                                     right = self._flopCount_subtree)

        self.constant = 0

    def visit_If(self, node, offset = 0):
        # Thread divergence
        self._ifCount += 1


    def visit_ID(self, node, offset = 0):
        symbol = self._lookUpAndGet(node)
        if symbol is None:
            # print " Symbol %s not found in ST "%(node.name)
            return
        if symbol in self._contableVars and is_complex_type(symbol.node):
            # print "Symbol " + str(symbol.name) + " is global "
            self._memAccess += 1
            self._currentAccess += 1
            self._memAccess_subtree = c_ast.BinaryOp( op = "+", 
                                                     left = c_ast.Constant(type = "integer", value = "1"),
                                                     right = self._memAccess_subtree)
        else:
            # print "Symbol : " + symbol.name + " not found in contableVars " + str([str(elem) for elem in self._contableVars])
            self._currentAccess = 0
            pass




